package com.xjai.controller;


import com.xjai.bean.ChatEntity;
import com.xjai.service.OllamaService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.stream.Collectors;

@Slf4j
@RestController
@RequestMapping("ollama")
public class OllamaController {

    @Resource
    private OllamaChatClient ollamaChatClient;

    @Resource
    private OllamaService ollamaService;

    @GetMapping("/ai/chat")
    public Object aiOllamaCall(@RequestParam String msg){
        log.info("prompt ------->> 111 :" + msg);
        return ollamaChatClient.call(msg);
    }

    @GetMapping("/ai/stream1")
    public Flux<ChatResponse> aiOllamaStream1(@RequestParam String msg){
        Prompt prompt = new Prompt(new UserMessage(msg));
        Flux<ChatResponse> stream = ollamaChatClient.stream(prompt);
        return stream;
    }


    @GetMapping("/ai/stream2")
    public List<String> aiOllamaStream2(@RequestParam String msg){
        Prompt prompt = new Prompt(new UserMessage(msg));
        Flux<ChatResponse> stream = ollamaChatClient.stream(prompt);

        List<String> collect = stream.toStream().map(chatResponse -> {
            String content = chatResponse.getResult().getOutput().getContent();
            log.info("content" + content);
            return content;
        }).collect(Collectors.toList());

        return collect;
    }

    @PostMapping("/ai/v3/doctor/stream")
    public void aiOllamaV3DoctorStream(@RequestBody ChatEntity chatEntity) {

        log.info(chatEntity.toString());
        String userName = chatEntity.getCurrentUserName();
        String message = chatEntity.getMessage();

        ollamaService.doDoctorStreamV3(userName, message);
    }

}